package com.wgz.ai.langchain4jai.memory;

import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import static dev.langchain4j.data.message.ChatMessageDeserializer.messagesFromJson;
import static dev.langchain4j.data.message.ChatMessageSerializer.messagesToJson;

/**
 * @author WGZ
 * @date 2025/3/5 16:23
 * 每个用户的持久化聊天记忆。
 */
public class ServiceWithPersistentMemoryForEachUserExample {
    private final static String API_KEY = "sk-svcacct-J-zoI616zoIN8ghM6amdG_SbFNc3vADd5yRa-53EQkr-lPN0tjD8QgWJfAmSuT3BlbkFJxBhbhopzngM8ZUB8KwByRFPnJ5bDUT1eq2Kop0OlggeqfzdsNSF45DQFZ94mgA";

    interface Assistant{
        String chat(@MemoryId int memoryId, @UserMessage String userMessage);
    }

    public static void main(String[] args) {
        ChatLanguageModel model = OpenAiChatModel.builder()
                .apiKey(API_KEY)
                .modelName("gpt-3.5-turbo")
                .logRequests(true)
                .logResponses(true)
                .build();

        ChatMemoryProvider chatMemory = memoryId ->MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(10)
                .chatMemoryStore(new PersistentChatMemoryStore())
                .build();

        Assistant assistant= AiServices.builder(Assistant.class)
                .chatMemoryProvider(chatMemory)
                .chatLanguageModel(model)
                .build();

        System.out.println(assistant.chat(1, "Hello, my name is Klaus"));
        System.out.println(assistant.chat(2, "Hi, my name is Francine"));
//        System.out.println(chat);
    }

    static class PersistentChatMemoryStore implements  ChatMemoryStore{


        private static final Map<String,String> MAP = new ConcurrentHashMap<>();

        @Override
        public List<ChatMessage> getMessages(Object o) {
            if (StringUtils.isEmpty(o)){
                return new ArrayList<>();
            }
            String json = MAP.get(o.toString());
            return messagesFromJson(json);
        }

        @Override
        public void updateMessages(Object o, List<ChatMessage> list) {
            if (StringUtils.isEmpty(o) || CollectionUtils.isEmpty(list)){
                return;
            }
            MAP.put(o.toString(),messagesToJson(list));
        }

        @Override
        public void deleteMessages(Object o) {
            if (StringUtils.isEmpty(o)){
                return;
            }
            MAP.remove(o.toString());
        }
    }
}
